import json
import uuid
from enum import Enum
from typing import Dict, Any

import websockets


class MCPTransportType(Enum):
    WEBSOCKET = "websocket"
    STDIO = "stdio"


class MCPClient:
    def __init__(self, transport_type: MCPTransportType = MCPTransportType.WEBSOCKET,
                 endpoint: str = "ws://localhost:8000/ws"):
        self.transport_type = transport_type
        self.endpoint = endpoint
        self.websocket = None

    async def connect(self):
        if self.transport_type == MCPTransportType.WEBSOCKET:
            endpoint_with_token = f"{self.endpoint}?token=my-secret-token"
            self.websocket = await websockets.connect(endpoint_with_token)
        else:
            raise NotImplementedError("STDIO transport not implemented yet")

    async def disconnect(self):
        if self.websocket:
            await self.websocket.close()

    async def list_tools(self) -> Dict[str, Any]:
        request = {
            "request_id": str(uuid.uuid4()),
            "action": "list_tools",
            "payload": None
        }
        await self.websocket.send(json.dumps(request))
        response = await self.websocket.recv()
        return json.loads(response)

    async def invoke_tool(self, tool_name: str, parameters: Dict[str, Any]) -> Any:
        request = {
            "request_id": str(uuid.uuid4()),
            "action": "invoke_tool",
            "payload": {
                "tool_name": tool_name,
                "parameters": parameters
            }
        }
        await self.websocket.send(json.dumps(request))
        response = await self.websocket.recv()
        response_data = json.loads(response)

        if not response_data.get('success'):
            raise Exception(response_data.get('error', 'Unknown error'))

        return response_data.get('data')

    async def __aenter__(self):
        await self.connect()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.disconnect()

    # 在MCPClient类中添加
    async def _stdio_communicate(self, request: dict) -> dict:
        import subprocess

        # 启动服务器进程（如果尚未运行）
        if not hasattr(self, 'server_process'):
            self.server_process = subprocess.Popen(
                ["python", "server/main.py"],
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                text=True
            )

        # 发送请求
        json_request = json.dumps(request)
        self.server_process.stdin.write(json_request + "\n")
        self.server_process.stdin.flush()

        # 读取响应
        response = self.server_process.stdout.readline()
        return json.loads(response)
